package me.tuyou.function;

import org.apache.flink.api.common.functions.MapPartitionFunction;
import org.apache.flink.api.java.ExecutionEnvironment;
import org.apache.flink.api.java.operators.DataSource;
import org.apache.flink.api.java.operators.MapPartitionOperator;
import org.apache.flink.api.java.operators.PartitionOperator;
import org.apache.flink.api.java.tuple.Tuple3;
import org.apache.flink.util.Collector;

public class MapPartitionTest {

    public static void main(String[] args) throws Exception {
        ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();
        DataSource<Tuple3<Integer, String, Double>> input = env.fromElements(
                Tuple3.of(1, "a", 1.0),
                Tuple3.of(1, "a", 1.1),
                Tuple3.of(2, "b", 2.2),
                Tuple3.of(2, "b", 2.3),
                Tuple3.of(3, "c", 3.3),
                Tuple3.of(3, "c", 3.4)
        );

        //分区
        PartitionOperator<Tuple3<Integer, String, Double>> partition = input.partitionByRange(1);
        //统计每个分区
        MapPartitionOperator<Tuple3<Integer, String, Double>, Integer> mapPartition = partition.mapPartition(new MapPartitionFunction<Tuple3<Integer, String, Double>, Integer>() {
            @Override
            public void mapPartition(Iterable<Tuple3<Integer, String, Double>> values, Collector<Integer> out) throws Exception {
                int count = 0;
                System.out.println("=================");
                for (Tuple3<Integer, String, Double> value : values) {
                    System.out.println(value);
                    count++;
                }
                if (count > 0) {
                    out.collect(count);
                }
            }
        });

        System.out.println("------------");
        mapPartition.print();
    }
}
